/*
 * Copyright (c) China Telecom Cloud Technology Co., Ltd. 2024-2025. All rights reserved.
 */
#ifndef KVCACHESDK_H
#define KVCACHESDK_H

#include <grpcpp/grpcpp.h>
#include <map>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <string>

#include "Common.h"
#include "MemoryPool.h"
#include "Segment.h"
#include "proto/kvcache.grpc.pb.h"

using grpc::Channel;
using grpc::ClientContext;
using grpc::Status;
using namespace kvcache;


/**
 *! @brief 共享内存文件对应的共享内存地址与大小
 * 
 */
struct SDKMapInfo {
    std::string path;
    uint8_t*    addr;
    uint64_t    len;

    SDKMapInfo() : path(""), addr(nullptr), len(0) {}
    
    SDKMapInfo(std::string path, uint8_t *addr, uint64_t len)
        : path(path), addr(addr), len(len) {}
};




class KVCacheSDK {
public:
    KVCacheSDK(const std::string &server_address, bool cuda)
        : _stub(KVCacheService::NewStub(grpc::CreateChannel(
              server_address, grpc::InsecureChannelCredentials()))),
          isCuda(cuda) {}
    

    /**
     *! @brief 主动unmap所有共享内存信息
     * 
     */
    ~KVCacheSDK();


    /**
     *! @brief 
     * 
     * @param hashIds 
     * @param tpNum 
     * @param doSave 
     * @return int32_t 
     */
    int32_t sealMemory(const std::vector<int64_t> &hashIds, 
                    int32_t tpNum,
                    bool doSave
                );


    /**
     *! @brief 
     * 
     * @param memSize 
     * @param seqs_blockHashIds 
     * @param startLayer 
     * @param endLayer 
     * @param layerNum 
     * @param layerCacheSize 
     * @param tpNum 
     * @return int32_t 
     */
    int32_t PrefetchPrefixCache(const int32_t memSize,
                            const MatrixT<BlockHashId> &seqs_blockHashIds,
                            const int32_t startLayer, 
                            const int32_t endLayer,
                            const int32_t layerNum, 
                            const int32_t layerCacheSize, 
                            int32_t tpNum
                        );
    
    
    /**
     *! @brief 批量查询前缀缓存
     * 
     * @param hashIdsList 
     * @return std::vector<std::vector<Segment>> 
     */
    std::vector<std::vector<Segment>> BatchQueryPrefixCache(
            std::vector<std::vector<int64_t>> hashIdsList);
    
    
    std::vector<Segment>
    AllocateMemory(const uint32_t memSize,
                   const std::vector<BlockHashId> &blockHashIds,
                   const int32_t writeInRecordSize, const int32_t startLayer,
                   const int32_t endLayer, const int32_t layerNum,
                   const int32_t layerCacheSize);
    
    
    /**
     *! @brief 
     * 
     * @param memSize 
     * @param type 
     */
    void LockMemoryPool(const uint32_t memSize, DeviceType type);
    
    
    /**
     *! @brief 
     * 
     * @return true 
     * @return false 
     */
    bool CheckIfKVClientReady();
    
    
    /**
     *! @brief 
     * 
     * @param memSize 
     */
    void CleanMemoryPool(const int32_t memSize);



    // 旧版本入参只包含hashId

    /**
     *! @brief 从存储预取特定hashIds的KV Cache到内存中(底层调用stub)
     * 
     * @param memSize         block_size个token的总KV Cache大小(字节)， 包含TP/PP
     * @param seqs_hashIds    待预取的hashId序列
     * @param startLayer      待预取的KV Cache的模型起始层(PP场景)
     * @param endLayer        待预取的KV Cache的模型结束层(PP场景)
     * @param layerNum        模型的总层数
     * @param layerCacheSize  单层block_size个token的总KV Cache大小(字节)， 包含TP/PP
     * @param tpNum 
     * 
     * @return int32_t 
     */
    int32_t PrefetchPrefixCache(const int32_t memSize,
                            const std::vector<std::vector<int64_t>> &seqs_hashIds,
                            const int32_t startLayer, 
                            const int32_t endLayer,
                            const int32_t layerNum, 
                            const int32_t layerCacheSize,
                            int32_t tpNum
                        );
    
    
    /**
     *! @brief 
     * 
     * @param memSize            内存块大小
     * @param hashIds            待分配Segment的块哈希ID
     * @param writeInRecordSize 
     * @param startLayer 
     * @param endLayer 
     * @param layerNum 
     * @param layerCacheSize 
     * @return std::vector<Segment> 
     */
    std::vector<Segment> AllocateMemory(const uint32_t memSize, 
                const std::vector<int64_t> &hashIds,
                const int32_t writeInRecordSize, 
                const int32_t startLayer,
                const int32_t endLayer, 
                const int32_t layerNum,
                const int32_t layerCacheSize);

private:
    /**
     *! @brief 检查特定大小的共享内存文件是否存在且与shmPath保持一致, 如果存在且文件一致则直
     *         接返回, 否则unmap原来的，并重新创建
     * 
     * @param memSize   共享内存大小
     * @param shmPath   预期的共享内存文件
     * 
     * @return int32_t  成功返回0否则返回错误码(可能会有异常抛出)
     */
    int32_t CheckAndCreateShm(uint32_t memSize, std::string shmPath);


    // 
    std::unique_ptr<KVCacheService::Stub> _stub;


    //! key: 共享内存大小
    // value: 对应的共享内存文件信息(路径，地址，长度)
    // std::map<uint32_t, SDKMapInfo> shmMapInfo;
    uInt32MapT<SDKMapInfo> shmMapInfo;


    bool isCuda;
};

namespace py = pybind11;

PYBIND11_MODULE(KVCacheSDK, m) {
    m.doc() = "pybind11 KVCacheSDK";
    py::class_<KVCacheSDK>(m, "KVCacheSDK")
        .def(py::init<const std::string &, bool>())
        .def("sealMemory", &KVCacheSDK::sealMemory)
        .def("PrefetchPrefixCache",
             py::overload_cast<const int32_t,
                               const std::vector<std::vector<int64_t>> &,
                               const int32_t, const int32_t, const int32_t,
                               const int32_t, int32_t>(
                 &KVCacheSDK::PrefetchPrefixCache))
        .def("PrefetchPrefixCache",
             py::overload_cast<const int32_t,
                               const std::vector<std::vector<BlockHashId>> &,
                               const int32_t, const int32_t, const int32_t,
                               const int32_t, int32_t>(
                 &KVCacheSDK::PrefetchPrefixCache))
        .def("AllocateMemory",
             py::overload_cast<const uint32_t, const std::vector<int64_t> &,
                               const int32_t, const int32_t, const int32_t,
                               const int32_t, const int32_t>(
                 &KVCacheSDK::AllocateMemory))
        .def("AllocateMemory",
             py::overload_cast<const uint32_t, const std::vector<BlockHashId> &,
                               const int32_t, const int32_t, const int32_t,
                               const int32_t, const int32_t>(
                 &KVCacheSDK::AllocateMemory))
        .def("LockMemoryPool", &KVCacheSDK::LockMemoryPool)
        .def("CheckIfKVClientReady", &KVCacheSDK::CheckIfKVClientReady)
        .def("CleanMemoryPool", &KVCacheSDK::CleanMemoryPool)
        .def("BatchQueryPrefixCache", &KVCacheSDK::BatchQueryPrefixCache);
    py::class_<Segment>(m, "Segment")
        .def("getSegmentId", &Segment::getSegmentId)
        .def("writeIn", &Segment::writeIn)
        .def("copyTo", &Segment::copyTo)
        .def("getIsFilled", &Segment::getIsFilled);
    pybind11::enum_<DeviceType>(m, "DeviceType")
        .value("CPU", DeviceType::DEVICE_TYPE_CPU)
        .value("GPU", DeviceType::DEVICE_TYPE_GPU)
        .value("NPU", DeviceType::DEVICE_TYPE_NPU)
        .export_values();
    pybind11::class_<BlockHashId>(m, "BlockHashId")
        .def(py::init<int64_t, int64_t, uint64_t, char *>())
        .def(py::init<>())
        .def_readwrite("hashId", &BlockHashId::hashId)
        .def_readwrite("parentHashId", &BlockHashId::parentHashId)
        .def_readwrite("localHashId", &BlockHashId::localHashId)
        .def_readonly("verifyHashId", &BlockHashId::verifyHashId);
}

#endif